"""模型管理模块，用于下载和验证模型"""

import subprocess
import psutil

from .utils.display import display
from .runner import OllamaService
from .utils.constants import MODEL_SIZES


class ModelManager:
    """管理模型下载和验证"""

    def __init__(self, config):
        """
        功能：初始化模型管理器
        :param config: 应用的配置信息
        :return: 无
        """
        self.config = config
        self.service = OllamaService()

    def model_exists(self, model_tag):
        """
        功能：检查模型是否已存在于本地
        :param model_tag: 要检查的模型标签
        :return: 如果模型存在返回True，否则返回False
        """
        cmd = ["ollama", "list"]
        try:
            # 执行ollama list命令并捕获输出
            result = subprocess.run(cmd, capture_output=True, text=True, check=False)
            return model_tag in result.stdout
        except (OSError, subprocess.SubprocessError) as e:
            display(f"⚠️ 模型存在检查失败: {e}")
            return False

    def ensure_model(self, model_tag):
        """
        功能：确保模型存在，若缺失则下载
        :param model_tag: 要确保可用的模型标签
        :return: 无（成功执行后模型可用）
        """
        self.service.start()
        display(f"检查模型 {model_tag} 是否已存在...")

        if self.model_exists(model_tag):
            display(f"✅ 模型 {model_tag} 已存在")
            return

        if not self.check_disk_space(model_tag):
            # 使用公共常量获取模型大小
            required = MODEL_SIZES.get(model_tag, 5) * 1.2
            display(f"🛑🛑🛑🛑 磁盘空间不足，无法下载 {model_tag} 模型")
            raise OSError(f"磁盘空间不足，至少需要 {required}GB 空间")

        display(f"🚀🚀 下载模型: {model_tag} (可能需要几分钟)...")
        try:
            with subprocess.Popen(
                ["ollama", "pull", model_tag],
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
            ) as process:
                while True:
                    output = process.stdout.readline()
                    if not output and process.poll() is not None:
                        break
                    if output and not (
                        output.startswith("pulling manifest")
                        or output.startswith("verifying sha256 digest")
                    ):
                        display(output.strip())

            if process.returncode == 0:
                display(f"✅ 模型 {model_tag} 下载完成")
            else:
                display(f"🛑🛑🛑 模型下载失败，退出码: {process.returncode}")
                raise RuntimeError(f"模型 {model_tag} 下载失败")
        except (OSError, subprocess.SubprocessError) as e:
            display(f"❌❌ 模型下载出错: {e}")
            raise

    def check_disk_space(self, model_tag):
        """
        功能：检查磁盘空间是否充足
        :param model_tag: 要下载的模型标签
        :return: 如果磁盘空间充足返回True，否则返回False
        """
        try:
            # 使用公共常量获取模型需要空间
            required_space = MODEL_SIZES.get(model_tag, 5)
            disk = psutil.disk_usage("/")
            free_gb = disk.free / (1024**3)
            # 要求额外保留20%的空间余量
            return free_gb >= required_space * 1.2
        except OSError as e:
            # 磁盘操作异常时提示并返回空间不足
            display(f"⚠️ 磁盘空间检查失败: {e}")
            return False
